#!/usr/bin/python
#-*- coding: utf-8 -*-
import random, string, array, binascii
from time import clock as now
__author__ = "lyricconch@gmail.com"
__version__ = "1.0b"

UNIT= "".join(map(chr,xrange(256)))
KEY_XOR= 0
KEY_ROT= 1
KEY_BOX= 2
"""
def base64key(input, type=KEY_XOR, DUP=False):
    if type == KEY_XOR:
        if DUP: return binascii.b2a_base64(array.array("i",input).tostring())
        return list(array.array("i", binascii.a2b_base64(input)))
    elif type == KEY_ROT:
        if DUP: return binascii.b2a_base64(array.array("B",input).tostring())
        return list(array.array("B", binascii.a2b_base64(input)))
    elif type == KEY_BOX:
        if DUP: return binascii.b2a_base64(input)
        return binascii.a2b_base64(input)
"""
def randomkey(len, type=KEY_XOR):
    random.seed();
    assert len<=256, "key length should be less than 256"
    if type == KEY_XOR:
        #return random.randrange(1<<len*8,(1<<len*8+1)-1)
        #return array.array("i",[random.randint(-1<<31,~(-1<<31)) for _ in xrange(len/4))]).tostring()
        key= list(UNIT)
        random.shuffle(key)
        return "".join(key)[:len]
    elif type == KEY_ROT:
        key= list(UNIT[:len])
        random.shuffle(key)
        return "".join(key)
    elif type == KEY_BOX: 
        key= list(UNIT)        
        random.shuffle(key)
        return "".join(key)
        
def invertkey(key, type=0):
    if type == KEY_XOR:
        return key    
    elif type == KEY_ROT: 
        return string.maketrans(key, UNIT[:len(key)])[:len(key)]
    elif type == KEY_BOX:  
        return string.maketrans(key,UNIT)

def transXor(data, xor): 
    key= array.array("i", xor if len(xor)%4==0 else xor*4)
    cnt= len(key)
    rem= len(data)%4  
    temp= array.array("i", data[:-rem] if rem else data)
    for j in xrange(cnt):
        cur= key[j]
        for i in xrange(j,len(temp),cnt):
            temp[i]^= cur
    if rem==0: return temp.tostring()
    return temp.tostring()+"".join([
        chr(ord(i)^ord(j)) 
        for i,j in zip(data[-rem:],xor)])        ;"""
    ardata= itertools.imap(int.__xor__,ardata,itertools.cycle(keyXor))
    temp= array.array("i",ardata).tostring() 
    #temp=struct.pack("%di"%(len(data)/4),*ardata)   
    return temp+data[-rem:] if rem else temp #"""

"""
def transXor2(data, keyXor):
    temp= map(ord,data) 
    temp= itertools.imap(int.__xor__,temp,itertools.cycle(keyXor))
    temp= map(chr,temp)
    
def transXor_(data, keyXor):    
    cnt2= keyXor.bit_length()/8*2
    data= binascii.b2a_hex(data)
    rem2= len(data)%cnt2    
    temp= [hex(long(data[i:i+cnt2],16)^keyXor)[3:-1]
        for i in xrange(0,len(data)-rem2,cnt2)]
    if rem2!=0: temp.append(
        hex(long(data[-rem2:],16)^keyXor)[-1-rem2:-1])    
    return binascii.a2b_hex("".join(temp)) #"""
    
def transRot(data, rot):  
    key= bytearray(rot)
    cnt= len(key) 
    rem= len(data)%cnt
    fix= len(data)-rem
    temp= bytearray(len(data))
    for i in xrange(cnt):
        temp[i:fix:cnt]= data[key[i]:fix:cnt]    
    if rem==0: return str(temp);
    for i in xrange(0,rem):
        o= key[i]; 
        while o>=rem: o=key[o]; 
        temp[fix+o]= data[fix+i]   
    return str(temp)    
    
def transBox(data, box):
    return data.translate(box)    
    
def transform(data, XOR=None, ROT=None, BOX=None, DEC=False): 
    def trace(when, TS=[0,0]): 
        """; TS[0],TS[1]= TS[1],now();
        print "%s%#.5s: %.60s%s"%(when+(TS[0]and"="or"@"),
            TS[1]-TS[0], binascii.b2a_hex(data), 30<len(data)and"..."or"")  #""";
        return True; 
    if DEC:
        assert trace("tran")
        if BOX: data= transBox(data, BOX); assert trace("dbox");
        if ROT: data= transRot(data, ROT); assert trace("drot");
        if XOR: data= transXor(data, XOR); assert trace("dxor");
    else:
        assert trace("data")
        if XOR: data= transXor(data, XOR); assert trace("exor");
        if ROT: data= transRot(data, ROT); assert trace("erot");
        if BOX: data= transBox(data, BOX); assert trace("ebox");
    return data
    
def testFAST(DS=10,XS=32,RS=80):
    CHAR= UNIT[32:127]
    t0= now()
    data= "".join([random.choice(CHAR) for _ in xrange(1001)])
    data= data*(DS/1001)+data[:DS%1001]
    t1= now()
    kxor,krot,kbox= randomkey(0,0,XS), randomkey(0,1,RS), randomkey(0,2)
    ixor,irot,ibox= invertkey(kxor,0), invertkey(krot,1), invertkey(kbox,2)    
    t2= now()
    tran= transform(data,kxor,krot,kbox,0)
    t3= now();
    undo= transform(tran,ixor,irot,ibox,1)
    t4= now();
    print "init: Bdata=%d, Cxor/Crot=%d/%d, Time=%.3f"%(
       len(data), len(kxor)*4, len(krot), t2-t1)
    print "done: Tenc=%.3f Tdec=%.3f"%(t3-t2, t4-t3)   
    print "data: %.60s%s\nfast: %.60s%s"%(
        data, 60<len(data)and"..."or"", 
        binascii.b2a_hex(tran), 30<len(tran)and"..."or"")
    assert data==undo

if __name__=="__main__":
    testFAST(1000000,100,32)





